import torch
import torch.nn as nn
from torch_scatter import scatter_mean, scatter_add
from einops import rearrange

import torch
import numpy as np
import torch.nn as nn
from timm.models.layers import trunc_normal_
from einops import rearrange, repeat
from torch_scatter import scatter_softmax, scatter

from torch_geometric.nn.pool import global_add_pool,global_mean_pool
from utils.normalization import Normalizer

ACTIVATION = {
    "gelu": nn.GELU,
    "tanh": nn.Tanh,
    "sigmoid": nn.Sigmoid,
    "relu": nn.ReLU,
    "leaky_relu": nn.LeakyReLU(0.1),
    "softplus": nn.Softplus,
    "ELU": nn.ELU,
    "silu": nn.SiLU,
}


class Graph_Physics_Attention_1D(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, slice_num=64):
        super().__init__()
        inner_dim = dim_head * heads
        self.dim_head = dim_head
        self.heads = heads
        self.scale = dim_head**-0.5
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

        self.temperature = nn.Parameter(torch.ones([1, heads, 1, 1]) * 0.5)
        self.graph_temperature = nn.Parameter(torch.ones([1, heads, 1]) * 0.5)

        self.in_project_x = nn.Linear(dim, inner_dim)
        self.in_project_fx = nn.Linear(dim, inner_dim)
        self.in_project_slice = nn.Linear(dim_head, slice_num)
        for l in [self.in_project_slice]:
            torch.nn.init.orthogonal_(l.weight)  # use a principled initialization
        self.to_q = nn.Linear(dim_head, dim_head, bias=False)
        self.to_k = nn.Linear(dim_head, dim_head, bias=False)
        self.to_v = nn.Linear(dim_head, dim_head, bias=False)
        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))

    def graph_forward(self, x, batch):
        # x: [B*N, C], batch: [B*N, 1]
        total_nodes = x.size(0)
        batch_size = batch.max().item() + 1

        ### (1) Slice
        fx_mid = self.in_project_fx(x).view(
            total_nodes, self.heads, self.dim_head
        )  # [B*N, H, D]
        x_mid = self.in_project_x(x).view(
            total_nodes, self.heads, self.dim_head
        )  # [B*N, H, D]

        slice_weights = self.softmax(
            self.in_project_slice(x_mid) / self.graph_temperature
        )  # [B*N, H, G]
        slice_norm = scatter_add(slice_weights, batch, dim=0)  # [B, H, G]

        slice_token = scatter_add(
            fx_mid.unsqueeze(-2) * slice_weights.unsqueeze(-1), batch, dim=0
        )  # [B, H, G, C]
        slice_token = slice_token / (slice_norm.unsqueeze(-1) + 1e-5)
        slice_token_1 = slice_token.clone()

        ### (2) Attention among slice tokens
        q_slice_token = self.to_q(slice_token)
        k_slice_token = self.to_k(slice_token)
        v_slice_token = self.to_v(slice_token)
        dots = torch.matmul(q_slice_token, k_slice_token.transpose(-1, -2)) * self.scale
        attn = self.softmax(dots)
        # attn = self.dropout(attn)
        out_slice_token = torch.matmul(attn, v_slice_token)  # B H G D
        slice_token_2 = out_slice_token.clone()

        ### (3) Deslice
        out_slice_token_expanded = out_slice_token[batch.squeeze()]  # [B*N, H, G, D]
        slice_weights_expanded = slice_weights  # [B*N, H, G]

        out_x = torch.sum(
            out_slice_token_expanded * slice_weights_expanded.unsqueeze(-1), dim=-2
        )  # [B*N, H, D]
        out_x_g = out_x.clone()
        out_x = rearrange(out_x, "n h d -> n (h d)")

        return self.to_out(out_x), slice_token_1, slice_token_2, out_x_g


class MLP(nn.Module):
    def __init__(self, n_input, n_hidden, n_output, n_layers=1, act="gelu", res=True):
        super(MLP, self).__init__()

        if act in ACTIVATION.keys():
            act = ACTIVATION[act]
        else:
            raise NotImplementedError
        self.n_input = n_input
        self.n_hidden = n_hidden
        self.n_output = n_output
        self.n_layers = n_layers
        self.res = res
        self.linear_pre = nn.Sequential(nn.Linear(n_input, n_hidden), act())
        self.linear_post = nn.Linear(n_hidden, n_output)
        self.linears = nn.ModuleList(
            [
                nn.Sequential(nn.Linear(n_hidden, n_hidden), act())
                for _ in range(n_layers)
            ]
        )

    def forward(self, x):
        x = self.linear_pre(x)
        for i in range(self.n_layers):
            if self.res:
                x = self.linears[i](x) + x
            else:
                x = self.linears[i](x)
        x = self.linear_post(x)
        return x


class Transolver_block(nn.Module):
    """Transformer encoder block."""

    def __init__(
        self,
        num_heads: int,
        hidden_dim: int,
        dropout: float,
        act="gelu",
        mlp_ratio=4,
        last_layer=False,
        slice_num=32,
    ):
        super().__init__()
        self.last_layer = last_layer
        self.ln_1 = nn.LayerNorm(hidden_dim)
        self.Attn = Graph_Physics_Attention_1D(
            hidden_dim,
            heads=num_heads,
            dim_head=hidden_dim // num_heads,
            dropout=dropout,
            slice_num=slice_num,
        )
        self.ln_2 = nn.LayerNorm(hidden_dim)
        self.mlp = MLP(
            hidden_dim,
            hidden_dim * mlp_ratio,
            hidden_dim,
            n_layers=0,
            res=False,
            act=act,
        )

    def forward(self, fx, batch):
        fx = self.Attn.graph_forward(self.ln_1(fx), batch) + fx
        fx = self.mlp(self.ln_2(fx)) + fx
        return fx
    


class Transolver_block(nn.Module):
    """Transformer encoder block."""

    def __init__(
        self,
        num_heads: int,
        hidden_dim: int,
        dropout: float,
        act="gelu",
        mlp_ratio=4,

        slice_num=32,
    ):
        super().__init__()

        self.ln_1 = nn.LayerNorm(hidden_dim)
        self.Attn = Graph_Physics_Attention_1D(
            hidden_dim,
            heads=num_heads,
            dim_head=hidden_dim // num_heads,
            dropout=dropout,
            slice_num=slice_num,
        )
        self.ln_2 = nn.LayerNorm(hidden_dim)
        self.mlp = MLP(
            hidden_dim,
            hidden_dim * mlp_ratio,
            hidden_dim,
            n_layers=0,
            res=False,
            act=act,
        )

    def forward(self, fx, batch):
        fx = self.Attn.graph_forward(self.ln_1(fx), batch)[0] + fx
        fx = self.mlp(self.ln_2(fx)) + fx
        return fx
    


class TransolverNet(nn.Module):
    def __init__(
        self,
        in_dim=7,
        n_layers=5,
        n_hidden=256,
        dropout=0,
        n_head=8,
        act="gelu",
        mlp_ratio=1,
        slice_num=32,
        params=None,
    ):
        super(TransolverNet, self).__init__()
        

        self.n_hidden = n_hidden
        try:
            satistic_times = params.dataset_size // params.batch_size
        except:
            satistic_times = 500


        self.normalizer = Normalizer(in_dim, satistic_times ,filter_zero=True)
        self.encoder = MLP(
            n_input=in_dim, 
            n_hidden=2*n_hidden, 
            n_output=n_hidden, 
            n_layers=1, 
            res=False
            )
        self.blocks = nn.ModuleList(
            [
                Transolver_block(
                    num_heads=n_head,
                    hidden_dim=n_hidden,
                    dropout=dropout,
                    act=act,
                    mlp_ratio=mlp_ratio,
                    slice_num=slice_num,
                )
                for _ in range(n_layers)
            ]
        )
        self.out_norm = nn.LayerNorm(n_hidden)
        

        self.initialize_weights()

    def initialize_weights(self):
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)


    def forward(self, x, batch, params=None):
        x = self.normalizer(x)
        x = self.encoder(x)
        for block in self.blocks:
            x = block(x, batch)
        x = self.out_norm(x)
        return x



if __name__ == '__main__':
    B, N = 3, 4
    dim = 7
    x = torch.randn(B, N, dim).cuda() 
    batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]).cuda()
    x_graph = x.reshape(-1, dim).cuda() 
    
    transo_deco = TransolverNet(7, 5, 256,dropout=0.0,n_head=8,act='gelu', slice_num=2).cuda()
    
    cd = transo_deco(x.reshape(3*4, -1), batch)
    print(cd.shape)